#! /usr/bin/env python
# coding=utf-8


# 对一个fastq文件进行分片 放到一个临时路径中

from Bio import SeqIO
from argparse import RawTextHelpFormatter

import os
import sys
import argparse
import shutil
import time

parser = argparse.ArgumentParser(
    description='''
    将fasta/fastq 文件进行分片
    用法：
    split_seq_to_a_tmp_file.py -i test.fq -s 10000 -o tmp_fq
    ''',formatter_class=RawTextHelpFormatter)



parser.add_argument('-i',
                help='fastq文件')

parser.add_argument('-s',
                help='分片大小 默认为10000')

parser.add_argument('-o',
                help='必须给定，输出文件')


# 获取参数
args = parser.parse_args()



if not args.i or not args.o:

    parser.print_help()

    sys.exit()

if not args.s:
    size = 10000
else:
    size = int(args.s)


infile = args.i 

tmp_file = args.o

try:
    shutil.rmtree(tmp_file)
except:
    pass

try:
    os.mkdir(tmp_file)
except:
    pass


if infile.split('.')[-1] in ['fastq','fq']:
    seq_type = 'fastq'
elif infile.split('.')[-1] in ['fasta','fa']:
    seq_type = 'fasta'
else:
    print('请给定fastq 或 fastq 文件')

n = 0 
k = 0
outfile = open(tmp_file+'/batch_'+str(k)+'.'+seq_type,'w')

for i in SeqIO.parse(infile,seq_type):

    SeqIO.write(i,outfile,seq_type)
    if n > size:
        print(k)
        outfile.close()
        n = 0
        k += 1
        outfile = open(tmp_file+'/batch_'+str(k)+'.'+seq_type,'w')
    n += 1

try:
    outfile.close()
except:
    pass
print('finished')